package com.shortlink.common.interceptor;

import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;

import java.util.Set;

/**
 * 滑动窗口限流
 */
public class LimitWindows {

    // 访问限制
    private final static long limit = 100;
    // 窗口时长（毫秒）
    private final static long windowDuration = 60 * 1000;

    private final static String LIMIT_WINDOW_LUA=
            "local window_start_time = ARGV[1] -ARGV[3]*1000 " +
            " redis.call('ZREMRANGEBYSCORE',KEYS[1],'-inf',window_start_time) " +
            " local now_request = redis.call('ZCARD',KEYS[1]) " +
            " if now_request < tonumber(ARGV[2]) then " +
            " redis.call('ZADD',KEYS[1],ARGV[1],ARGV[1]) " +
            "     return 1 " +
            "else " +
            "      return 0 " +
            " end ";

    public boolean allowRequest(StringRedisTemplate stringRedisTemplate, String key) {
        // 获取当前时间
        long currentTime = System.currentTimeMillis();
        // 计算窗口开始时间
        long windowStart = currentTime - windowDuration;
        // 删除窗口开始时间之前的所有数据
        stringRedisTemplate.opsForZSet().removeRangeByScore(key, 0, windowStart);
        // 使用 ZCARD 计算当前窗口内的请求数
        Long current = stringRedisTemplate.opsForZSet().zCard(key);

        // 判断是否允许访问
        if (current < limit) {
            // 如果窗口内请求数未达到限制，添加当前请求
            stringRedisTemplate.opsForZSet().add(key, String.valueOf(currentTime), currentTime);
            return true;
        }

        return false;
    }
}